using EnzymeCore
import EnzymeCore.EnzymeRules: forward

f(x) = x^2

function forward(::Const{typeof(f)}, RT::Type{<:DuplicatedNoNeed}, x::Duplicated)
    return 10+2*x.val*x.dval
end

function forward(::Const{typeof(f)}, RT::Type{<:BatchDuplicatedNoNeed}, x::BatchDuplicated{T, N}) where {T, N}
    return NTuple{N, T}(1000+2*x.val*dv for dv in x.dval)
end

function forward(::Const{typeof(f)}, RT::Type{<:Duplicated}, x::Duplicated)
    return Duplicated(func.val(x.val), 100+2*x.val*x.dval)
end

function forward(::Const{typeof(f)}, RT::Type{<:BatchDuplicated}, x::BatchDuplicated{T, N}) where {T,N}
    return BatchDuplicated(func.val(x.val), NTuple{N, T}(10000+2*x.val*dv for dv in x.dval))
end

import EnzymeCore: Annotation
function has_frule(@nospecialize(f))
    TT = Tuple{<:Annotation{Core.Typeof(f)}, Type{<:Annotation}, Vararg{<:Annotation}}
    isapplicable(forward, TT)
end

# Do we need this one?
function has_frule(@nospecialize(f), @nospecialize(RT))
    TT = Tuple{<:Annotation{Core.Typeof(f)}, Type{RT}, Vararg{<:Annotation}}
    isapplicable(forward, TT)
end

# Do we need this one?
function has_frule(@nospecialize(f), @nospecialize(RT), @nospecialize(TT))
    TT = Base.unwrap_unionall(TT)
    TT = Tuple{<:Annotation{Core.Typeof(f)}, Type{RT}, TT.parameters...}
    isapplicable(forward, TT)
end

# Base.hasmethod is a precise match we want the broader query.
function isapplicable(@nospecialize(f), @nospecialize(TT); world=Base.get_world_counter())
    tt = Base.to_tuple_type(TT)
    sig = Base.signature_type(f, tt)
    return !isempty(Base._methods_by_ftype(sig, -1, world)) # TODO cheaper way of querying?
end

@assert has_frule(f)
@assert has_frule(f, Duplicated)
@assert has_frule(f, DuplicatedNoNeed)
@assert has_frule(f, BatchDuplicated)
@assert has_frule(f, BatchDuplicatedNoNeed)
@assert has_frule(f, Duplicated, Tuple{<:Duplicated})
@assert has_frule(f, DuplicatedNoNeed, Tuple{<:Duplicated})
@assert has_frule(f, BatchDuplicated, Tuple{<:BatchDuplicated})
@assert has_frule(f, BatchDuplicatedNoNeed, Tuple{<:BatchDuplicated})

@assert !has_frule(f, Duplicated, Tuple{<:BatchDuplicated})
@assert !has_frule(f, DuplicatedNoNeed, Tuple{<:BatchDuplicated})
@assert !has_frule(f, BatchDuplicated, Tuple{<:Duplicated})
@assert !has_frule(f, BatchDuplicatedNoNeed, Tuple{<:Duplicated})